import pandas as pd

import json
from zhipuai import ZhipuAI
import time

api_key = "e07484c365a3d8b468a4d9e06b36317a.momrwta4EftapAwx"

client = ZhipuAI(api_key=api_key)

with open('prompt/prompt_getitem.txt') as f:
    prompt = f.read()


#发送请求
def send_req(tar_text):
    while True:
        response = client.chat.asyncCompletions.create(
            model="glm-3-turbo",
            messages=[{"role": "user", "content": prompt % tar_text}],
            top_p=0.7,
            temperature=0.9,
        )
        print(response)
        if response.task_status == 'PROCESSING':
            return response.id
        if response.task_status == 'FAILED':
            return None    
        else:
            print('wait')
        time.sleep(10)

#查询异步请求状态
def query_status(task_id):
    response = client.chat.asyncCompletions.retrieve_completion_result(id=task_id)
    
    if response.task_status == 'SUCCESS':
        return {'status': 1, 'content': response.choices[0].message.content}
    if response.task_status == 'FAILED':
        return {'status': -1}
    return {'status':0}
    

def check_wait(wait_list): 
    tmp_list = []
    for item in wait_list:
        tmp_res = query_status(item[0])
        if tmp_res['status'] == 0:
            tmp_list.append(item)
        
        if tmp_res['status'] == 1:
            print(tmp_res['content'])
            
        
        if tmp_res['status'] == -1:
            print('trans error ')
        time.sleep(2)
    return tmp_list 






WAIT_NUM = 3
wait_list = []


for i in range(1,3):
    print(i)
    #print(row)
    with open('test_data/news%d.txt' % i) as f:
        tar_content = f.read()
    #print(cur_item)
    tmp_id = send_req(tar_content)
    
    
    if tmp_id:
        #cur_item['task_id'] = tmp_id
        wait_list.append([tmp_id,tar_content])
    else:
        print('trans error')
    #if last_check_ts == 0:
        #last_check_ts = time.time()
    time.sleep(3)

    while len(wait_list) >= WAIT_NUM:
    #if time.time() - last_check_ts > CHECK_INTERVAL:
        print('check_list %d' % len(wait_list))
        time.sleep(5)
        wait_list = check_wait(wait_list)
        #last_check_ts = time.time()
    
print('final check ----------')

while len(wait_list) > 0:
    print('check_list %d' % len(wait_list))
    wait_list = check_wait(wait_list)
    time.sleep(5)   
    
